#_*_coding:utf-8_*_
import tensorflow as tf
# from tensorflow.examples.tutorials.mnist import input_data
import tensorflow.contrib.slim as slim
from collections import OrderedDict
import torch

from econet import ECONet
# from econetv1 import ECONet

import numpy as np


def load_pt_model(path):
    pt_model = torch.load(path)['state_dict']

    pt_model_inception = OrderedDict()
    pt_model_resnet = OrderedDict()
    pt_model_fc = OrderedDict()

    for k, v in pt_model.items():
        if 'num_batches_tracked' not in k:
            if 'res' in k:
                pt_model_resnet[k] = v
            elif 'fc' in k:
                pt_model_fc[k] = v
            else:
                pt_model_inception[k] = v

    # Update for matching tensorflow variable scopes
    pt_model_inception.update(pt_model_resnet)
    pt_model_inception.update(pt_model_fc)

    print(len(pt_model_inception.keys()))
    
    return pt_model_inception


def save_param(pt_model, save_path):
    # Hyperparameters
    opt = {
        'weight_decay': 0.0, 
        'net2d_keep_prob': 0.5,
        'net3d_keep_prob': 0.5,
        'num_segments': 3,
        'num_classes': 400 
    }

    # torch.load()
    # Input data
    with tf.name_scope('input'):
        x = tf.placeholder(tf.float32, [None, 224, 224, 3], name = 'x_input')
        y = tf.placeholder(tf.float32, [None, 10], name = 'y_input')

    logits, _ = ECONet(x, opt=opt)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        # glb_vars = tf.global_variables() # shows every variable being used.
        # for var in
        vars_list = tf.global_variables()

        i = 0
        for (k, v), variable in zip(pt_model.items(), vars_list):

            if len(v.size()) == 4:
                v = v.permute(2, 3, 1, 0).contiguous()
            elif len(v.size()) == 5:
                v = v.permute(2, 3, 4, 1, 0).contiguous()
            elif len(v.size()) == 2:
                v = v.permute(1, 0).contiguous()


            v_np = v.cpu().numpy()

            assert v_np.shape == variable.shape, \
                'Assigned variable shape {} not equals to original variable shape {}'.format(v_np.shape, variable.shape)
            
            # print(variable.shape, variable.name)
            sess.run(variable.assign(v_np))
            # print(variable.eval())
            # i += 1 
            # break

        saver = tf.train.Saver()
        saver.save(sess, save_path)   


def compare_param(pt_model, read_path):
    # Hyperparameters
    opt = {
        'weight_decay': 0.0, 
        'net2d_keep_prob': 0.5,
        'net3d_keep_prob': 0.5,
        'num_segments': 3,
        'num_classes': 400 
    }

    # Input data
    with tf.name_scope('input'):
        x = tf.placeholder(tf.float32, [None, 224, 224, 3], name = 'x_input')
        y = tf.placeholder(tf.float32, [None, 10], name = 'y_input')

    logits, _ = ECONet(x, opt=opt)

    with tf.Session() as sess:
        saver = tf.train.Saver() 
        sess.run(tf.global_variables_initializer())
        saver.restore(sess, read_path)
        vars_list = tf.global_variables()

        # print(vars_list)
        i = 0
        for (k, v), variable in zip(pt_model.items(), vars_list):

            if len(v.size()) == 4:
                v = v.permute(2, 3, 1, 0).contiguous()
            elif len(v.size()) == 5:
                v = v.permute(2, 3, 4, 1, 0).contiguous()
            elif len(v.size()) == 2:
                v = v.permute(1, 0).contiguous()


            v_np = v.cpu().numpy()

            # print(i)
            if not np.array_equal(v_np, variable.eval()):
                print("The variable {} in tensorflow checkpoint is different from torch's".format(variable.name))

            i += 1 
            # sess.run(variable.assign(a))
            # print()


if __name__ == '__main__':
    pt_model_path = '/home/paper99/Projects/ECO-pytorch/ECO_Full_rgb_model_Kinetics.pth.tar'
    tf_model_path = 'experiments/test_model/ECOfull_kinetics.ckpt'
    
    pt_model = load_pt_model(pt_model_path)
    save_param(pt_model, tf_model_path)
    # compare_param(pt_model, tf_model_path)
    # test_graph()
    # calc_param()
    # save_param()

